import os
# import openslide
import tqdm
import time
from PIL import Image
import numpy as np
import cv2
import copy
import xml.etree.ElementTree as ET
import shutil
import argparse
import threading
import torch
import torchvision
OPENSLIDE_PATH = r'C:\openslide-win64-20231011\bin'




import os
if hasattr(os, 'add_dll_directory'):
    # Windows
    with os.add_dll_directory(OPENSLIDE_PATH):
        import openslide
else:
    import openslide


# python3 automated_segmentaition_ndpi.py --ndpi_path '/media/alex/FA5EB5A15EB556DB1/uterus/2022.3.15子宫确诊'
parser = argparse.ArgumentParser(description='Ndpi to xml')
parser.add_argument('--ndpi_path', '-f', type=str, default=False, help='Load ndpi from a .ndpi file')
args = parser.parse_args()

ndpi_dir = args.ndpi_path

window_size = 640

for ndpi_path in os.listdir(ndpi_dir):
    if (ndpi_path[-4:]=='ndpi'):
        start_pred_time = time.time()
        ndpi_path = os.path.join(ndpi_dir,ndpi_path)
        png_path = './png_test' # png dir ,cropped by ndpi



        if os.path.exists(png_path):
            shutil.rmtree(png_path )
            os.mkdir(png_path )
        else:
            os.mkdir(png_path )
        time1 = time.time()
        slide = openslide.open_slide(ndpi_path)
        print ('opentime :',time.time()-time1)
        width, height = slide.dimensions
        print (width,height)
        index_x = 0
        index_y = 0
        index = 0


        def process_and_save_image(region,index_x, index_y,index,png_path):
            # simg1=os.path.join(png_path,'{}_0000.png'.format(str(index).zfill(4)))
            simg1=os.path.join(png_path,'{}_{}_{}_0000.png'.format(index_x,index_y,str(index).zfill(4)))
            region = region.convert('RGB')
            region.save(simg1)
        # crop ndpi ,get png images
        threads = []
        while True:
            while True:
                tile_size_x = window_size if index_x+window_size<width else width-index_x+1
                tile_size_y = window_size if index_y+window_size<height else height-index_y+1

                start_read_time = time.time()
                region = slide.read_region((index_x, index_y), 0, (tile_size_x, tile_size_y))
                print('read_region_time',time.time()-start_read_time)
                t = threading.Thread(target=process_and_save_image,args=(region,index_x, index_y,index,png_path))
                t.start()
                threads.append(t)
                
                index +=1
                if index_x+window_size<width:
                    index_x +=window_size-100 
                else:
                    index_x = 0
                    break

            if index_y+window_size<height:
                    index_y +=window_size-100  
            else:
                index_y = 0
                break
        for t in threads:
            t.join()
        slide.close()




        from ultralytics import YOLO

        # 加载一个模型
        model = YOLO(r'runs/detect/yolov8-l/weights/best.pt')  # 从YAML建立并转移权重

        results = model.predict(png_path,imgsz=window_size,conf=0.5,stream=True)

        # from ultralytics.engine.results

        # load xml 
        xml_path = './module.xml'

        tree = ET.parse(xml_path)
        root = tree.getroot()
        annotation =  list(root.iter('Annotation'))[0]

        # draw cells
        regions  = list(annotation.iter('Regions'))[0]
        region  = list(annotation.iter('Region'))[0]
        Vertices = list(region.iter('Vertices'))[0]
        vertex = list(Vertices.iter('Vertex'))[0]

        root.remove(annotation)
        annotation.remove(regions)
        regions.remove(region)
        region.remove(Vertices)
        Vertices.remove(vertex)

        ID = 1
        scores = []
        points_all = []
        for result in results:
            if len(result.boxes)>0:
                image_x = int(os.path.split(result.path)[-1].split('_')[0])
                image_y = int(os.path.split(result.path)[-1].split('_')[1])
                for score in result.boxes.conf:
                    scores.append(score)
                for xywh in result.boxes.xywh:
                    x,y,w,h = int(xywh[0]-xywh[2]/2)+image_x,int(xywh[1]-xywh[3]/2)+image_y,int(xywh[2]),int(xywh[3])
                    # points_xyxy = [
                    #     [x,y],[x+w,y],[x+w,y+h],[x,y+h],
                    #     # [x,y-1]
                    # ]
                    points_xyxy = [x,y,x+w,y+h]
                    points_all .append(copy.deepcopy(points_xyxy))
        print (f'bbox num before nms : {len(points_all)}')
        points_all = torch.tensor(points_all).float()
        scores = torch.tensor(scores).float()
        classes = torch.zeros((len(points_all))).float()
        points_keep_index = torchvision.ops.batched_nms(points_all,scores, classes,iou_threshold=0.2)
        print (f'bbox num after nms : {len(points_keep_index)}')


        for points_index in list(points_keep_index):
            point_index = points_index.item()
            points_xyxy = points_all[points_index]
            region_to_save = copy.deepcopy(region)
            Vertices_to_save = copy.deepcopy(Vertices )

            region_to_save.set('Id',str(ID))
            region_to_save.set('Text',str(scores[points_index].item())[:5])
            region_to_save.set('DisplayId',str(ID))
            x1,y1,x2,y2 = int(points_xyxy[0].item()),int(points_xyxy[1].item()),int(points_xyxy[2].item()),int(points_xyxy[3].item())
            points = [
                [x1,y1],[x2,y1],[x2,y2],[x1,y2]
            ]
            for point in points:
                vertex_to_save = copy.deepcopy(vertex)
                vertex_to_save.set('X',str(point[0]))
                vertex_to_save.set('Y',str(point[1]))
                Vertices_to_save.append(vertex_to_save)
            region_to_save.append(Vertices_to_save)
            regions.append(region_to_save)
            ID += 1

        annotation.append(regions)
        root.append(annotation)

        tree.write(ndpi_path[:-5]+'.xml')

        print (f'{ndpi_path} total_time',time.time()-start_pred_time)
